import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import warnings
from matplotlib.cbook import MatplotlibDeprecationWarning
# 禁用警告
pd.options.mode.chained_assignment = None
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=MatplotlibDeprecationWarning)

# 设置中文字体
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

# 初始化pro接口
pro = ts.pro_api('a5acdd993403f9e5ae13a16f2d9182b8e35d67e3beddfebd4bb43712')

# 获取股票历史数据
def get_stock_data(ts_code, start_date, end_date):
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df = df.sort_values(by='trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    numeric_columns = df.select_dtypes(include=[np.number]).columns
    df[numeric_columns] = df[numeric_columns].astype(float)
    df['k_data'] = df.apply(
        lambda x: [
            x.name.strftime('%Y-%m-%d'),  # 修改日期格式为 '-' 分隔（适配Matplotlib）
            x['open'], x['close'],
            x['low'], x['high']
        ],
        axis=1
    )
    df['ts_code'] = ts_code  # 添加股票代码列，用于图表标题
    return df

# 移动平均线交叉策略
def moving_average_crossover_strategy(df, short_window=5, long_window=20):
    df['short_ma'] = df['close'].rolling(window=short_window).mean()
    df['long_ma'] = df['close'].rolling(window=long_window).mean()
    df['signal'] = 0
    df.loc[df['short_ma'] > df['long_ma'], 'signal'] = 1
    df.loc[df['short_ma'] < df['long_ma'], 'signal'] = -1
    df['position'] = df['signal'].diff()
    print("移动平均线交叉策略")
    return df

# 基于机器学习的量化策略
def machine_learning_strategy(df):
    features = ['open', 'close', 'high', 'low', 'vol', 'amount']
    df['target'] = df['close'].shift(-1) > df['close']
    df['target'] = df['target'].astype(int)
    df.dropna(inplace=True)
    X = df[features]
    y = df['target']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print("基于机器学习策略")
    print(f"机器学习模型准确率：{accuracy:.4f}")
    df['signal'] = model.predict(df[features])
    df['position'] = df['signal'].diff()
    return df

# 应用止损和止盈规则
def apply_stop_loss_take_profit(df, stop_loss=0.05, take_profit=0.1):
    df['stop_loss'] = df['close'] * (1 - stop_loss)
    df['take_profit'] = df['close'] * (1 + take_profit)
    df['position'] = 0
    position = 0
    for i in range(1, len(df)):
        if df['signal'].iloc[i] == 1 and position == 0:
            position = 1
            buy_price = df['close'].iloc[i]
        elif position == 1:
            current_price = df['close'].iloc[i]
            if (current_price / buy_price - 1) <= -stop_loss or (current_price / buy_price - 1) >= take_profit:
                position = 0
        df['position'].iloc[i] = position
    return df

# 回测策略
def backtest_strategy(df):
    df['returns'] = df['close'].pct_change()
    df['strategy_returns'] = df['position'].shift(1) * df['returns']
    df['cumulative_returns'] = (1 + df['strategy_returns']).cumprod()
    return df

# 计算夏普比率
def calculate_sharpe_ratio(df, risk_free_rate=0.03):
    returns = df['strategy_returns'].dropna()
    sharpe_ratio = (returns.mean() - risk_free_rate) / returns.std()
    return sharpe_ratio

# 生成ECharts格式K线数据（简化，直接使用DataFrame数据）
def generate_echarts_kline_data(df):
    return {
        'dates': df.index.strftime('%Y-%m-%d').tolist(),  # 统一日期格式
        'open': df['open'].tolist(),
        'close': df['close'].tolist(),
        'low': df['low'].tolist(),
        'high': df['high'].tolist(),
        'ma5': df['close'].rolling(5).mean().tolist(),
        'ma10': df['close'].rolling(10).mean().tolist(),
        'ma20': df['close'].rolling(20).mean().tolist(),
        'ma30': df['close'].rolling(30).mean().tolist()
    }

# 可视化结果（关键修改点：子图创建和日期格式处理）
def visualize_results(df, strategy_name):
    ts_code = df['ts_code'].iloc[0]
    kline_data = generate_echarts_kline_data(df)
    
    # 创建包含两个子图的figure（主图：K线+信号，子图：收益率）
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 12))  # 解包子图对象
    
    # 主图：绘制K线
    ax1.set_title(f'{ts_code} {strategy_name} 策略效果')
    ax1.set_xlabel('日期')
    ax1.set_ylabel('价格')
    
    # 绘制K线柱
    for i in range(len(kline_data['dates'])):
        open_price = kline_data['open'][i]
        close_price = kline_data['close'][i]
        low = kline_data['low'][i]
        high = kline_data['high'][i]
        color = '#00ff00' if close_price > open_price else '#ff0000'
        ax1.vlines(kline_data['dates'][i], low, high, color='k', lw=0.7)
        ax1.bar(kline_data['dates'][i], close_price - open_price, bottom=open_price, width=0.6, color=color)
    
    # 绘制均线
    ax1.plot(kline_data['dates'], kline_data['ma5'], label='MA5', color='blue', linestyle='--', alpha=0.7)
    ax1.plot(kline_data['dates'], kline_data['ma10'], label='MA10', color='orange', linestyle='--', alpha=0.7)
    ax1.plot(kline_data['dates'], kline_data['ma20'], label='MA20', color='green', linestyle='--', alpha=0.7)
    ax1.plot(kline_data['dates'], kline_data['ma30'], label='MA30', color='red', linestyle='--', alpha=0.7)
    
    # 绘制买卖信号（转换日期为字符串）
    buy_signals = df[df['position'] == 1]
    sell_signals = df[df['position'] == -1]
    
    ax1.scatter(
        buy_signals.index.strftime('%Y-%m-%d'),  # 直接使用索引的字符串表示
        buy_signals['close'],
        marker='^',
        color='g',
        label='买入信号',
        zorder=2
    )
    ax1.scatter(
        sell_signals.index.strftime('%Y-%m-%d'),
        sell_signals['close'],
        marker='v',
        color='r',
        label='卖出信号',
        zorder=2
    )
    
    ax1.legend()
    ax1.grid(True, linestyle='--', alpha=0.5)
    plt.setp(ax1.get_xticklabels(), rotation=45, ha='right')  # 旋转x轴标签
    
    # 子图：累计收益率
    ax2.plot(df.index, df['cumulative_returns'], label=f'累计收益率', color='purple', linestyle='-')
    ax2.set_title('策略累计收益率')
    ax2.set_xlabel('日期')
    ax2.set_ylabel('累计收益率（%）')
    ax2.grid(True, linestyle='--', alpha=0.5)
    ax2.legend()
    
    plt.tight_layout()  # 调整子图间距
    plt.show()

# 主程序
if __name__ == "__main__":
    ts_codes = ["688128.SH", "600278.SH"]  # 示例股票代码
    start_date = '20250101'  # 调整起始日期（避免未来数据）
    end_date = '20250501'
    
    strategy_names = ["移动平均线交叉", "机器学习"]
    strategy_functions = [moving_average_crossover_strategy, machine_learning_strategy]
    
    for ts_code in ts_codes:
        print(f"处理股票代码：{ts_code}...")
        df = get_stock_data(ts_code, start_date, end_date)
        
        for strategy_name, strategy_func in zip(strategy_names, strategy_functions):
            print(f"  应用策略：{strategy_name}")
            strategy_df = strategy_func(df.copy())
            strategy_df = apply_stop_loss_take_profit(strategy_df)
            strategy_df = backtest_strategy(strategy_df)
            
            # 计算指标
            cumulative_return = strategy_df['cumulative_returns'].iloc[-1] - 1
            sharpe_ratio = calculate_sharpe_ratio(strategy_df)
            
            print(f"    累计收益率：{cumulative_return:.4f}")
            print(f"    夏普比率：{sharpe_ratio:.4f}")
            
            visualize_results(strategy_df, strategy_name)